import os
import cv2
import numpy as np
from paddleocr import PaddleOCR
from functools import wraps
import time
from fuzzywuzzy import fuzz


class OCRProcessor:
    def __init__(self):
        # 初始化PaddleOCR模型，支持中文识别
        self.ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False)
    
    def recognize_text(self, image_path):
        """
        识别图像中的文本
        
        Args:
            image_path: 图像文件路径
            
        Returns:
            results: 识别结果列表，每项包含文本和置信度
            annotated_image: 标注了识别结果的图像
        """
        # 确保文件存在
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图像文件不存在: {image_path}")
        
        # OCR识别
        result = self.ocr.ocr(image_path, cls=True)
        
        # 提取结果
        ocr_results = []
        if result and len(result) > 0 and result[0]:
            for line in result[0]:
                text = line[1][0]  # 文本内容
                confidence = line[1][1]  # 置信度
                box = line[0]  # 文本框坐标
                ocr_results.append({
                    'text': text,
                    'confidence': float(confidence),
                    'box': box
                })
        
        # 标注图像
        annotated_image = self._annotate_image(image_path, result)
        
        return ocr_results, annotated_image
    
    def _annotate_image(self, image_path, result):
        """在图像上标注识别结果"""
        # 读取原始图像
        img = cv2.imread(image_path)
        
        if result and len(result) > 0 and result[0]:
            for line in result[0]:
                # 获取文本框坐标
                points = line[0]
                points = np.array(points).astype(np.int32)
                
                # 绘制边框
                cv2.polylines(img, [points], True, (0, 255, 0), 2)
                
                # 添加文本
                text = line[1][0]
                cv2.putText(img, text, (points[0][0], points[0][1] - 10), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
        
        return img
    
    def preprocess_image(self, image_path):
        """
        图像预处理以提高OCR准确率
        
        Args:
            image_path: 原始图像路径
            
        Returns:
            processed_path: 处理后的图像路径
        """
        # 读取图像
        img = cv2.imread(image_path)
        
        # 转为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # 应用自适应阈值二值化
        binary = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, 11, 2
        )
        
        # 应用中值滤波去噪
        denoised = cv2.medianBlur(binary, 3)
        
        # 保存处理后的图像
        processed_path = image_path.replace('.', '_processed.')
        cv2.imwrite(processed_path, denoised)
        
        return processed_path 

# def fuzzy_match_nutrition(text, threshold=80):
#     best_match = None
#     best_score = 0
#
#     for db_name, info in NUTRITION_DB.items():
#         # 使用比率匹配算法
#         score = fuzz.ratio(text.lower(), db_name.lower())
#         if score > threshold and score > best_score:
#             best_score = score
#             best_match = (db_name, info)
#
#     return best_match